# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).

# --------------------------------------------------------
# CroCo model during pretraining
# --------------------------------------------------------

import torch
import torch.nn as nn
torch.backends.cuda.matmul.allow_tf32 = True # for gpu >= Ampere and pytorch >= 1.12
from functools import partial

from models.blocks import Block, DecoderBlock, PatchEmbed
from models.pos_embed import get_2d_sincos_pos_embed, RoPE2D
from models.masking import RandomMask

class CroCoNet(nn.Module):
    def __init__(self,
                 img_size=224,           # input image size，图像尺寸
                 patch_size=16,          # patch_size，补丁尺寸
                 mask_ratio=0.9,         # ratios of masked tokens，交叉视图补全预训练掩码比例
                 enc_embed_dim=768,      # encoder feature dimension，编码器嵌入维度
                 enc_depth=12,           # encoder depth，编码器深度(编码器块重复次数)
                 enc_num_heads=12,       # encoder number of heads in the transformer block，编码器头数量
                 dec_embed_dim=512,      # decoder feature dimension，解码器嵌入维度
                 dec_depth=8,            # decoder depth，解码器深度(解码器块重复次数)
                 dec_num_heads=16,       # decoder number of heads in the transformer block，解码器头数量
                 mlp_ratio=4,            # MLP全连接层隐藏层维度的扩展比例
                 norm_layer=partial(nn.LayerNorm, eps=1e-6), # 层归一化
                 norm_im2_in_dec=True,   # whether to apply normalization of the 'memory' = (second image) in the decoder，是否对token特征2做层归一化
                 pos_embed='cosine',     # positional embedding (either cosine or RoPE100)，位置编码类型
                ):

        super(CroCoNet, self).__init__()

        # patch embeddings  (with initialization done as in MAE)，补丁嵌入(掩码自编码器MAE初始化)
        self.patch_embed = PatchEmbed(img_size, patch_size, 3, enc_embed_dim) # 输入图像尺寸，补丁尺寸，输入通道，编码器嵌入维度

        # mask generations，设置掩码生成器
        self._set_mask_generator(self.patch_embed.num_patches, mask_ratio) # 输入补丁数量，掩码比例

        self.pos_embed = pos_embed # 位置编码类型
        if pos_embed=='cosine': # 余弦位置编码
            # positional embedding of the encoder
            enc_pos_embed = get_2d_sincos_pos_embed(enc_embed_dim, int(self.patch_embed.num_patches**.5), n_cls_token=0)
            self.register_buffer('enc_pos_embed', torch.from_numpy(enc_pos_embed).float()) # 相当于加入self.enc_pos_embed
            # positional embedding of the decoder
            dec_pos_embed = get_2d_sincos_pos_embed(dec_embed_dim, int(self.patch_embed.num_patches**.5), n_cls_token=0)
            self.register_buffer('dec_pos_embed', torch.from_numpy(dec_pos_embed).float()) # 相当于加入self.dec_pos_embed
            # pos embedding in each block
            self.rope = None # nothing for cosine
        elif pos_embed.startswith('RoPE'): # eg RoPE100，旋转位置编码
            self.enc_pos_embed = None # nothing to add in the encoder with RoPE
            self.dec_pos_embed = None # nothing to add in the decoder with RoPE
            if RoPE2D is None: raise ImportError("Cannot find cuRoPE2D, please install it following the README instructions")
            freq = float(pos_embed[len('RoPE'):])
            self.rope = RoPE2D(freq=freq) # 旋转位置编码
        else:
            raise NotImplementedError('Unknown pos_embed '+pos_embed)

        # transformer for the encoder，编码器相关定义
        self.enc_depth = enc_depth # 编码器深度(编码器块重复次数)
        self.enc_embed_dim = enc_embed_dim # 编码器嵌入维度，默认768
        self.enc_blocks = nn.ModuleList([ # 编码器块，旋转位置编码
            Block(enc_embed_dim, enc_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, rope=self.rope) # 输入编码器嵌入维度，编码器头数量，MLP全连接层隐藏层维度的扩展比例，qkv矩阵偏置，层归一化，旋转位置编码
            for i in range(enc_depth)]) # 重复构造Block
        self.enc_norm = norm_layer(enc_embed_dim) # 层归一化

        # masked tokens，设置token掩码
        self._set_mask_token(dec_embed_dim) # 输入解码器嵌入维度

        # decoder，设置解码器
        self._set_decoder(enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec) # 输入编码器嵌入维度，解码器嵌入维度，解码器深度，解码器头数量，解码器深度，MLP全连接层隐藏层维度的扩展比例，层归一化，是否对token特征2做层归一化

        # prediction head，设置CroCo交叉视图补全预训练的输出头
        self._set_prediction_head(dec_embed_dim, patch_size) # 输入解码器嵌入维度，补丁尺寸

        # initializer weights，权重初始化
        self.initialize_weights()

    # 设置掩码生成器
    def _set_mask_generator(self, num_patches, mask_ratio):
        self.mask_generator = RandomMask(num_patches, mask_ratio) # 随机生成掩码，输入补丁数量，掩码比例

    # 设置token掩码
    def _set_mask_token(self, dec_embed_dim):
        self.mask_token = nn.Parameter(torch.zeros(1, 1, dec_embed_dim)) # 创建零张量，(1,1,解码器嵌入维度)

    # 设置解码器
    def _set_decoder(self, enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec):
        self.dec_depth = dec_depth # 解码器深度(解码器块重复次数)
        self.dec_embed_dim = dec_embed_dim # 解码器嵌入维度
        # transfer from encoder to decoder
        self.decoder_embed = nn.Linear(enc_embed_dim, dec_embed_dim, bias=True) # 线性层，输入编码器嵌入维度(768)，输出解码器嵌入维度(512)
        # transformer for the decoder
        self.dec_blocks = nn.ModuleList([ # 解码器块，旋转位置编码
            DecoderBlock(dec_embed_dim, dec_num_heads, mlp_ratio=mlp_ratio, qkv_bias=True, norm_layer=norm_layer, norm_mem=norm_im2_in_dec, rope=self.rope) # 输入解码器嵌入维度，解码器头数量，MLP全连接层隐藏层维度的扩展比例，qkv矩阵偏置，层归一化，是否对token特征2做层归一化，旋转位置编码
            for i in range(dec_depth)]) # 重复构造Block
        # final norm layer
        self.dec_norm = norm_layer(dec_embed_dim) # 层归一化

    # 设置CroCo交叉视图补全预训练的输出头
    def _set_prediction_head(self, dec_embed_dim, patch_size):
         self.prediction_head = nn.Linear(dec_embed_dim, patch_size**2 * 3, bias=True) # 线性层，输入解码器嵌入维度，输出补丁尺寸^2*3(映射为补丁RGB)，(B,N,补丁尺寸^2*3)

    # 初始化权重
    def initialize_weights(self):
        # patch embed
        self.patch_embed._init_weights()
        # mask tokens
        if self.mask_token is not None: torch.nn.init.normal_(self.mask_token, std=.02)
        # linears and layer norms
        self.apply(self._init_weights)

    # 初始化权重
    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            # we use xavier_uniform following official JAX ViT:
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    # ViT编码器
    def _encode_image(self, image, do_mask=False, return_all_blocks=False):
        """
        image has B x 3 x img_size x img_size
        do_mask: whether to perform masking or not
        return_all_blocks: if True, return the features at the end of every block
                           instead of just the features from the last block (eg for some prediction heads)
        """
        # embed the image into patches  (x has size B x Npatches x C)
        # and get position if each return patch (pos has size B x Npatches x 2)
        x, pos = self.patch_embed(image) # 补丁嵌入，输入图像，输出token和补丁位置编码
        # add positional embedding without cls token，不加cls_token
        if self.enc_pos_embed is not None: # 旋转位置编码不运行
            x = x + self.enc_pos_embed[None,...] # 位置嵌入
        # apply masking
        B,N,C = x.size()
        if do_mask: # 掩码输入图像，用于预训练
            masks = self.mask_generator(x) # 随机生成token掩码(掩盖补丁)
            x = x[~masks].view(B, -1, C) # 对token掩码
            posvis = pos[~masks].view(B, -1, 2) # 对补丁位置编码掩码
        else: # 不进行掩码
            B,N,C = x.size()
            masks = torch.zeros((B,N), dtype=bool) # token掩码为零张量，(B,N)
            posvis = pos
        # now apply the transformer encoder and normalization
        if return_all_blocks: # 返回所有块的输出
            out = [] # out中包含了每个编码器块的输出结果
            for blk in self.enc_blocks: # 编码器块，旋转位置编码
                x = blk(x, posvis) # 输入token和补丁位置编码，输出token特征(旋转位置编码嵌入)
                out.append(x)
            out[-1] = self.enc_norm(out[-1]) # 对编码器最后输出做层归一化
            return out, pos, masks # 返回所有块输出，补丁位置编码，token掩码(或零张量)
        else:
            for blk in self.enc_blocks:
                x = blk(x, posvis)
            x = self.enc_norm(x)
            return x, pos, masks

    # Transformer解码器
    def _decoder(self, feat1, pos1, masks1, feat2, pos2, return_all_blocks=False):
        """
        return_all_blocks: if True, return the features at the end of every block
                           instead of just the features from the last block (eg for some prediction heads)

        masks1 can be None => assume image1 fully visible
        """
        # encoder to decoder layer，嵌入维度转换
        visf1 = self.decoder_embed(feat1) # 线性层，编码器嵌入维度(768) -> 解码器嵌入维度(512)
        f2 = self.decoder_embed(feat2)
        # append masked tokens to the sequence
        B,Nenc,C = visf1.size()
        if masks1 is None: # downstreams，下游任务
            f1_ = visf1 # 解码器嵌入维度的token特征1
        else: # pretraining，预训练
            Ntotal = masks1.size(1)
            f1_ = self.mask_token.repeat(B, Ntotal, 1).to(dtype=visf1.dtype)
            f1_[~masks1] = visf1.view(B * Nenc, C)
        # add positional embedding
        if self.dec_pos_embed is not None: # 旋转位置编码不运行
            f1_ = f1_ + self.dec_pos_embed # 位置嵌入
            f2 = f2 + self.dec_pos_embed
        # apply Transformer blocks
        out = f1_
        out2 = f2
        if return_all_blocks: # 返回所有块的输出
            _out, out = out, [] # out中包含了每个解码器块的输出结果
            for blk in self.dec_blocks: # 解码器块，旋转位置编码
                _out, out2 = blk(_out, out2, pos1, pos2) # 输入token特征1，token特征2，补丁位置编码1，补丁位置编码2
                out.append(_out)
            out[-1] = self.dec_norm(out[-1]) # 对解码器最后输出做层归一化
        else:
            for blk in self.dec_blocks:
                out, out2 = blk(out, out2, pos1, pos2)
            out = self.dec_norm(out)
        return out

    # 图像按照补丁词向量化，作为标签
    def patchify(self, imgs):
        """
        imgs: (B, 3, H, W)
        x: (B, L, patch_size**2 *3)
        """
        p = self.patch_embed.patch_size[0] # 补丁尺寸
        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0

        h = w = imgs.shape[2] // p # 高宽上的补丁数量
        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p)) # 调整形状，(B,3,高补丁数量,补丁尺寸,宽补丁数量,补丁尺寸)
        x = torch.einsum('nchpwq->nhwpqc', x) # 调整形状，(B,高补丁数量,宽补丁数量,补丁尺寸,补丁尺寸,3)
        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3)) # 调整形状为词向量，嵌入维度为补丁RGB，(B,补丁数量,补丁尺寸^2*3)

        return x

    def unpatchify(self, x, channels=3):
        """
        x: (N, L, patch_size**2 *channels)
        imgs: (N, 3, H, W)
        """
        patch_size = self.patch_embed.patch_size[0]
        h = w = int(x.shape[1]**.5)
        assert h * w == x.shape[1]
        x = x.reshape(shape=(x.shape[0], h, w, patch_size, patch_size, channels))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], channels, h * patch_size, h * patch_size))
        return imgs

    def forward(self, img1, img2):
        """
        img1: tensor of size B x 3 x img_size x img_size
        img2: tensor of size B x 3 x img_size x img_size

        out will be    B x N x (3*patch_size*patch_size)
        masks are also returned as B x N just in case
        """
        # encoder of the masked first image
        feat1, pos1, mask1 = self._encode_image(img1, do_mask=True) # ViT编码器，输入img1，进行掩码，输出token特征1，补丁位置编码1，token掩码1
        # encoder of the second image
        feat2, pos2, _ = self._encode_image(img2, do_mask=False) # 共享权重ViT编码器，输入img2，不掩码，输出token特征2，补丁位置编码2
        # decoder
        decfeat = self._decoder(feat1, pos1, mask1, feat2, pos2) # Transformer解码器，输入token特征1，补丁位置编码1，token掩码1，token特征2，补丁位置编码2
        # prediction head
        out = self.prediction_head(decfeat) # CroCo交叉视图补全预训练的输出头，使用线性层将解码器输出映射到补丁RGB，(B,N,补丁尺寸^2*3)
        # get target
        target = self.patchify(img1) # 将图像1按照补丁词向量化，作为标签
        return out, mask1, target
